use std::ops::Range;
use std::sync::Arc;

use datafusion::arrow::array::{ArrayBuilder, ArrayRef, ArrowPrimitiveType, Float64Array, Float64Builder, GenericStringArray, GenericStringBuilder, Int64Array, Int64Builder, PrimitiveBuilder, UInt32Array, UInt32Builder, UInt64Array, UInt64Builder};
use datafusion::arrow::buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer};
use datafusion::arrow::datatypes::{DataType, TimeUnit};

pub trait ArrayBuilderExt {
    fn slice(&self, data_type: &DataType, range: Range<usize>) -> ArrayRef;
}

macro_rules! primitive_slice {
    ($self: tt, $range: tt, $builder: ty, $array: ty) => {
        {
            let builder = unsafe {
                $self.as_any().downcast_ref_unchecked::<$builder>()
            };
            let values = &builder.values_slice()[$range.clone()];
            let validity = builder.validity_slice().map(|s|{
                let len = $range.end - $range.start;
                let start = $range.start / 8;
                let left = $range.start % 8;
                let end = (left + len).div_ceil(8);
                let ss = &s[start..end];
                NullBuffer::new(BooleanBuffer::new(
                    Buffer::from_slice_ref(ss), left, len,
                ))
            });
            let array = <$array>::new(
                Buffer::from_slice_ref(values).into(), validity,
            );
            Arc::new(array)
        }
    };
}

impl ArrayBuilderExt for Box<dyn ArrayBuilder> {
    #[inline]
    fn slice(&self, data_type: &DataType, range: Range<usize>) -> ArrayRef {
        boxed_slice(&self, data_type, range)
    }
}

impl ArrayBuilderExt for &Box<dyn ArrayBuilder> {
    #[inline]
    fn slice(&self, data_type: &DataType, range: Range<usize>) -> ArrayRef {
        boxed_slice(self, data_type, range)
    }
}

impl <T: ArrowPrimitiveType> ArrayBuilderExt for PrimitiveBuilder<T> {
    fn slice(&self, data_type: &DataType, range: Range<usize>) -> ArrayRef {
        slice_impl(self, data_type, range)
    }
}

// impl<T: ArrayBuilder> ArrayBuilderExt for &T {
//     fn slice(&self, data_type: &DataType, range: Range<usize>) -> ArrayRef {
//         slice::<T>(&self, data_type, range)
//     }
// }

fn boxed_slice(me: &Box<dyn ArrayBuilder>, data_type: &DataType, range: Range<usize>) -> ArrayRef {
    match data_type {
        DataType::Int64 => {
            slice_impl(
                unsafe { me.as_any().downcast_ref_unchecked::<Int64Builder>() },
                data_type,
                range,
            )
        }
        DataType::Float64 => {
            slice_impl(
                unsafe { me.as_any().downcast_ref_unchecked::<Float64Builder>() },
                data_type,
                range,
            )
        }
        DataType::UInt64 | DataType::Timestamp(TimeUnit::Millisecond, _) => {
            slice_impl(
                unsafe { me.as_any().downcast_ref_unchecked::<UInt64Builder>() },
                data_type,
                range,
            )
        }
        DataType::UInt32 => {
            slice_impl(
                unsafe { me.as_any().downcast_ref_unchecked::<UInt32Builder>() },
                data_type,
                range,
            )
        }
        DataType::Utf8 => {
            slice_impl(
                unsafe { me.as_any().downcast_ref_unchecked::<GenericStringBuilder<i32>>() },
                data_type,
                range,
            )
        }
        other => {
            unimplemented!("ArrayBuilderExt data_type: {} not implement", other)
        }
    }
}

fn slice_impl<T: ArrayBuilder>(me: &T, data_type: &DataType, range: Range<usize>) -> ArrayRef {
    match data_type {
        DataType::Int64 => {
            primitive_slice!(me, range, Int64Builder, Int64Array)
        }
        DataType::Float64 => {
            primitive_slice!(me, range, Float64Builder, Float64Array)
        }
        DataType::UInt64 | DataType::Timestamp(TimeUnit::Millisecond, _) => {
            primitive_slice!(me, range, UInt64Builder, UInt64Array)
        }
        DataType::UInt32 => {
            primitive_slice!(me, range, UInt32Builder, UInt32Array)
        }
        DataType::Utf8 => {
            let builder = unsafe {
                me.as_any().downcast_ref_unchecked::<GenericStringBuilder<i32>>()
            };
            let offset_range = range.start..range.end + 1;
            let offsets = &builder.offsets_slice()[offset_range];
            let value_range = (offsets[0] as usize)..(offsets[offsets.len() - 1] as usize);
            let offsets = OffsetBuffer::<i32>::new(
                Buffer::from_slice_ref(offsets).into()
            );
            let values = &builder.values_slice()[value_range];
            let values = Buffer::from_slice_ref(values);
            let validity = builder.validity_slice().map(|s| {
                let len = range.end - range.start;
                let start = range.start / 8;
                let left = range.start % 8;
                let end = (left + len).div_ceil(8);
                let ss = &s[start..end];
                NullBuffer::new(BooleanBuffer::new(
                    Buffer::from_slice_ref(ss), left, len,
                ))
            });
            let array = GenericStringArray::<i32>::new(
                offsets, values, validity,
            );
            Arc::new(array)
        }
        other => {
            unimplemented!("ArrayBuilderExt data_type: {} not implement", other)
        }
    }
}

#[cfg(test)]
mod test {
    use datafusion::arrow::array::{Array, ArrayBuilder, Float64Array, Float64Builder, GenericStringArray, GenericStringBuilder, Int64Array, Int64Builder};
    use datafusion::arrow::datatypes::DataType;

    use crate::array_builder_ext::ArrayBuilderExt;

    #[test]
    fn test_int64() {
        let mut builder = Int64Builder::with_capacity(100);
        builder.append_value(1);
        builder.append_value(2);
        builder.append_null();
        let b: Box<dyn ArrayBuilder> = Box::new(builder);
        let a = b.slice(&DataType::Int64, 0..3);
        assert_eq!(3, a.len());
        let a0 = unsafe {
            a.as_any().downcast_ref_unchecked::<Int64Array>()
        };
        assert_eq!(1, a0.value(0));
        assert_eq!(2, a0.value(1));
        assert!(a0.is_null(2));
    }

    #[test]
    fn test_f64() {
        let mut builder = Float64Builder::new();
        builder.append_value(1.);
        builder.append_value(2.);
        builder.append_null();
        let b: Box<dyn ArrayBuilder> = Box::new(builder);
        let a = b.slice(&DataType::Float64, 0..3);
        assert_eq!(3, a.len());
        let a0 = unsafe {
            a.as_any().downcast_ref_unchecked::<Float64Array>()
        };
        assert_eq!(1., a0.value(0));
        assert_eq!(2., a0.value(1));
        assert!(a0.is_null(2));
    }

    #[test]
    fn test_utf8() {
        let mut builder = GenericStringBuilder::<i32>::new();
        builder.append_value("123");
        builder.append_value("abcde");
        builder.append_null();
        let b: Box<dyn ArrayBuilder> = Box::new(builder);
        let a = b.slice(&DataType::Utf8, 0..3);
        assert_eq!(3, a.len());
        let a0 = unsafe {
            a.as_any().downcast_ref_unchecked::<GenericStringArray<i32>>()
        };
        assert_eq!("123", a0.value(0));
        assert_eq!("abcde", a0.value(1));
        assert!(a0.is_null(2));
    }
}